import os
import zipfile
from io import BytesIO


def to_dir(path: str):
    if path is None:
        return path

    path = path.strip()
    if len(path) == 0:
        return path

    if path[-1] in ('/', '\\'):
        path = path[:-1]
    return path


class FilePaths:
    def __init__(self, root):
        self.paths = get_paths_from_root(root)

    def __len__(self):
        return len(self.paths)

    def __iter__(self):
        return _Iter(self)

    @staticmethod
    def open(path: str, mode='r'):
        pos = path.find('||')
        if pos < 0:
            return open(path, mode)
        zf = zipfile.ZipFile(path[:pos])
        return zf.open(path[pos + 2:], 'r' if 'r' in mode else 'w')


class _Iter:
    def __init__(self, file_enum):
        self._file_enum = file_enum
        self._index = 0

    def __len__(self):
        return len(self._file_enum.paths)

    def __next__(self):
        if self._index >= len(self._file_enum.paths):
            raise StopIteration()
        path = self._file_enum.paths[self._index]
        self._index += 1
        return path


def get_paths_from_root(root: str):
    root = to_dir(root)
    result = []
    for name in os.listdir(root):
        path = root + os.sep + name
        if os.path.isdir(path):
            result.extend(get_paths_from_root(path))
        elif len(name) >= 4 and name[-4:].upper() == '.ZIP':
            with zipfile.ZipFile(path) as file:
                result.extend([path + '||' + name for name in get_paths_from_zip(file)])
        else:
            result.append(path)
    return result


def get_paths_from_zip(file: zipfile.ZipFile):
    return [info.filename for info in file.infolist() if not info.is_dir()]


if __name__ == '__main__':
    fe = FilePaths('../..')
    print(len(fe))
    # for path in fe:
    #     if r'\algorithm.zip' in path:
    #         print(path)
    #         with FileEnum.open(path, 'r') as f:
    #             print(f.readlines())
    # for name in get_paths_from_root('../..'):
    #     if r'\algorithm.zip' in name:
    #         print(name)
    # file = zipfile.ZipFile("../../algorithm.zip")
    # names = [info.filename for info in file.infolist() if not info.is_dir()]
    # print(names)
    # with file.open('algorithm/api/plc.py') as f:
    #     print(f.readlines())
    # file.close()

class ZipFileReader:
    def __init__(self, path_or_fp):
        self._z_file = zipfile.ZipFile(path_or_fp)
        self._index = 0

    def close(self):
        self._z_file.close()

    @staticmethod
    def from_path(path):
        return ZipFileReader(zipfile.ZipFile(path))

    @staticmethod
    def from_path(file):
        return ZipFileReader(zipfile.ZipFile(file))

    @staticmethod
    def from_bytes(bytes):
        f = BytesIO()
        f.write(bytes)
        return ZipFileReader(zipfile.ZipFile(f))

    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        self.close()

    def __iter__(self):
        self._index = 0
        return self

    def __next__(self):
        if self._index == self.num:
            self.z.close()
            raise StopIteration
        with self.z.open(self.file_paths[self._index], 'r') as f:
            img_bytes = f.read()
            img_array = cv2.imdecode(np.frombuffer(img_bytes, np.uint8), cv2.IMREAD_COLOR)

        self._index += 1
        return img_array

    def __len__(self):
        return self.num

